import numpy as np
import matplotlib.pyplot as plt
class PIDController:
    def __init__(self, Kp, Ki, Kd, setpoint):
        self.Kp = Kp
        self.Ki = Ki
        self.Kd = Kd
        self.setpoint = setpoint
        self.prev_error = 0
        self.integral = 0

    def compute(self, feedback):
        error = self.setpoint - feedback
        self.integral += error
        # derivative = error - self.prev_error
        derivative = (error - self.prev_error) / step_size

        output = self.Kp * error + self.Ki * self.integral + self.Kd * derivative

        self.prev_error = error
        return output
##################################################################
##################################################################
class Control:
    def __init__(self, desired_heading, desired_speed):
        Kp_heading = 270#270#300
        Ki_heading = 0
        Kd_heading = 1800#1800#1600
        Kp_speed = 310
        Ki_speed = 0.015 # 0.06
        Kd_speed = 0
        self.heading_controller = PIDController(Kp_heading, Ki_heading, Kd_heading, desired_heading)
        self.speed_controller = PIDController(Kp_speed, Ki_speed, Kd_speed, desired_speed)

    def update(self, current_heading, current_speed):
        heading_control = self.heading_controller.compute(current_heading)
        speed_control = self.speed_controller.compute(current_speed)
        return heading_control, speed_control
##################################################################
##################################################################
def speed(u, v, r, Fu):
    m11 = 710.56
    m22 = 1018.56
    d11 = 132.74
    d22 = 174.77
    u_dot = (m22 * v * r / m11) - (d11 * u / m11) + (Fu / m11)
    v_dot = -(m11 * u * r / m22) - (d22 * v / m22)
    return u_dot, v_dot

def angle(u, v, r, Tr):
    m11 = 710.56
    m22 = 1018.56
    m33 = 3918.4
    d33 = 35403.39
    r_dot = ((m11 - m22) * u * v / m33) - (d33 * r / m33) + (Tr / m33)
    return r_dot
##################################################################
##################################################################
# 设置期望的航向和航速
desired_heading = 90/57.3  # 期望的航向角度
desired_speed = 5  # 期望的航速

# 创建控制器实例
controller = Control(desired_heading, desired_speed)

# 迭代循环
step_size = 0.05
total_steps = 2001  # 循环次数
time = np.arange(0, (total_steps+1) * step_size, step_size)
# 模拟无人艇当前状态
current_heading = np.zeros(total_steps+1, dtype=float)  # 当前航向角度
current_speed = np.zeros(total_steps+1, dtype=float) # 当前航速
r = np.zeros(total_steps+1, dtype=float)
v = np.zeros(total_steps+1, dtype=float)
Fu = np.zeros(total_steps+1, dtype=float)
Tr = np.zeros(total_steps+1, dtype=float)

##################################################################
##################################################################
for i in range(1, total_steps+1):
    # 更新控制器并获取航向和航速控制信号
    heading_control, speed_control = controller.update(current_heading[i-1], current_speed[i-1])
    # print(speed_control)

    # 航向求解角加速度
    Tr[i] = heading_control * 81.7
    if Tr[i] < -2460:
        Tr[i] = -2460
    elif Tr[i] > 2460:
        Tr[i] = 2460
    else:
        Tr[i] = Tr[i]
    r_dot = angle(current_speed[i - 1], v[i - 1], r[i - 1], Tr[i])
    # r_dot = angle(current_heading[i-1], v, r, 0)
    # r_dot = angle(current_heading[i], v[i], r[i], Tr[i])

    # 航速求解两个加速度
    Fu[i] = speed_control * 90
    if Fu[i] < 0:
        Fu[i] = 0
    elif Fu[i] > 680:
        Fu[i] = 680
    else:
        Fu[i] = Fu[i]
    u_dot, v_dot = speed(current_speed[i-1], v[i-1], r[i-1], Fu[i])
    # u_dot, v_dot = speed(current_heading[i], v[i], r[i], Fu[i])
    # u_dot, v_dot = speed(current_heading[i-1], v, r, 680)
    # print(f'u_dot大小为{u_dot}')

    # 积分
    current_speed[i] = current_speed[i-1] + u_dot * step_size
    v[i] = v[i-1] + v_dot * step_size

    r[i] = r[i-1] + r_dot * step_size
    current_heading[i] = current_heading[i-1] + r[i-1] * step_size + r_dot * step_size * step_size / 2
##################################################################
##################################################################
# current_heading = [x * 57.3 for x in current_heading]
fig = plt.figure(1)
plt.plot(time, current_speed, ls = '-', c = 'r', lw = '1')
plt.plot(time, np.full((len(time)),desired_speed), ls = '--', c = 'b', lw = '1', alpha = 1)
plt.xlabel('Time/t')
plt.ylabel('U(m/s)')
# plt.gca().set_aspect('auto')

fig2 = plt.figure(2)
plt.plot(time, [x * 57.3 for x in current_heading], ls = '-', c = 'r', lw = '1')
# plt.plot(time, current_heading, ls = '-', c = 'r', lw = '2')
plt.plot(time, np.full((len(time)),desired_heading * 57.3), ls = '--', c = 'b', lw = '1', alpha = 1)
plt.xlabel('Time/t')
plt.ylabel('angle(°)')
# plt.gca().set_aspect('auto')

##################################################################
##################################################################
if len(np.where(abs(Tr) >= 2460+1e-5)[0]) > 0.5:
    raise Exception('扭矩有错误')
print(f'Tr的值为{Tr.tolist()}')


if (len(np.where((Fu - 680) >= 1e-5)[0]) > 0.5) or (len(np.where(Fu < 0)[0]) > 0.5):
    raise Exception('推力有错误')
print(f'Fu的值为{Fu.tolist()}')

print("Heading Control:", [i * 57.3 for i in current_heading])
print("Speed Control:", current_speed.tolist())

if len(np.where(abs(current_speed - desired_speed) <= 1e-4)[0]) != 0:
    t = time[np.where(abs(current_speed - desired_speed) <= 1e-4)[0][0]]
    print(f'时间{t}s时速度接近收敛')
else:
    print(f'航速在{total_steps * step_size}s内没有收敛')

if len(np.where(abs(current_heading - desired_heading) <= 1e-6)[0]) != 0:
    t = time[np.where(abs(current_heading - desired_heading) <= 1e-5)[0][0]]
    print(f'{t}s时航向接近收敛')
else:
    print(f'航向在{total_steps * step_size}s内没有收敛')

plt.show()